import torch
import argparse
from mmengine.config import Config
from mmengine.registry import init_default_scope
import numpy as np
import subprocess
import sys


def parse_args():
    parser = argparse.ArgumentParser(description='Get DiffusionDet flops')
    parser.add_argument('config', help='config file path')
    parser.add_argument('--shape', nargs='+', type=int, default=[3, 640, 640],
                        help='input image size')
    args = parser.parse_args()
    return args


def format_number(num):
    """格式化数字，添加适当的单位 (K, M, G, T)"""
    if num < 1e3:
        return f"{num:.2f}"
    elif num < 1e6:
        return f"{num / 1e3:.2f}K"
    elif num < 1e9:
        return f"{num / 1e6:.2f}M"
    elif num < 1e12:
        return f"{num / 1e9:.2f}G"
    else:
        return f"{num / 1e12:.2f}T"


def main():
    args = parse_args()

    # 定义分隔线
    split_line = '=' * 60

    # 加载配置
    cfg = Config.fromfile(args.config)

    # 初始化默认作用域
    init_default_scope(cfg.get('default_scope', 'mmdet'))

    try:
        # 导入所需模块
        from mmdet.registry import MODELS

        # 构建模型
        model = MODELS.build(cfg.model)
        if torch.cuda.is_available():
            model = model.cuda()
        model.eval()

        # 准备输入
        if len(args.shape) == 3:
            input_shape = tuple(args.shape)
        else:
            input_shape = (3, args.shape[0], args.shape[1])

        batch_size = 1
        dummy_input = torch.randn((batch_size, *input_shape))
        if torch.cuda.is_available():
            dummy_input = dummy_input.cuda()

        # 计算参数量
        total_params = sum(p.numel() for p in model.parameters())

        # 创建一个临时的包装类来计算 FLOPs
        class ModelWrapper(torch.nn.Module):
            def __init__(self, model):
                super().__init__()
                self.model = model

            def forward(self, x):
                # 仅使用骨干网络和颈部网络进行前向传播
                # 这样可以避免处理复杂的检测头和其他组件
                feats = self.model.backbone(x)
                if hasattr(self.model, 'neck'):
                    feats = self.model.neck(feats)
                return feats

        # 创建包装模型
        wrapped_model = ModelWrapper(model)

        # 尝试使用 fvcore 计算 FLOPs
        try:
            # 尝试导入 fvcore，如果不存在则安装
            try:
                from fvcore.nn import FlopCountAnalysis
            except ImportError:
                print("正在安装 fvcore 库...")
                subprocess.check_call([sys.executable, "-m", "pip", "install", "fvcore"])
                from fvcore.nn import FlopCountAnalysis

            # 使用包装模型计算 FLOPs
            flops_counter = FlopCountAnalysis(wrapped_model, dummy_input)
            flops = flops_counter.total()

            # 打印结果
            print(f"{split_line}")
            print(f"Input shape: {input_shape}")
            print(f"FLOPs (骨干网络+颈部网络): {format_number(flops)} ({flops:,})")
            print(f"Params: {format_number(total_params)} ({total_params:,})")
            print(f"{split_line}")

            # 尝试获取更详细的 FLOPs 分布
            print("\n各组件 FLOPs 分布:")
            print(f"{'-' * 60}")
            flops_table = flops_counter.by_module()
            for module_name, module_flops in sorted(flops_table.items(), key=lambda x: x[1], reverse=True)[
                                             :10]:  # 只显示前10个
                percentage = 100.0 * module_flops / flops
                print(f"{module_name:<40} {format_number(module_flops):<15} {percentage:.2f}%")
            print(f"{'-' * 60}")

        except Exception as flops_error:
            print(f"使用 fvcore 计算 FLOPs 失败: {flops_error}")

            # 尝试使用 thop 计算 FLOPs
            try:
                # 如果没有安装 thop，先安装它
                try:
                    import thop
                except ImportError:
                    print("正在安装 thop 库...")
                    subprocess.check_call([sys.executable, "-m", "pip", "install", "thop"])
                    import thop

                # 使用 thop 计算 MACs
                from thop import profile

                # 计算骨干网络的 MACs
                backbone_macs, _ = profile(
                    model.backbone,
                    inputs=(dummy_input,),
                    verbose=False
                )
                backbone_flops = 2 * backbone_macs

                # 初始化总 FLOPs
                total_flops = backbone_flops

                # 打印骨干网络 FLOPs
                print(f"{split_line}")
                print(f"Input shape: {input_shape}")
                print(f"骨干网络 FLOPs: {format_number(backbone_flops)} ({backbone_flops:,})")

                # 如果有颈部网络，尝试计算其 FLOPs
                if hasattr(model, 'neck'):
                    try:
                        # 获取骨干网络输出
                        with torch.no_grad():
                            backbone_output = model.backbone(dummy_input)

                        # 计算颈部网络 FLOPs
                        neck_macs, _ = profile(
                            model.neck,
                            inputs=(backbone_output,),
                            verbose=False
                        )
                        neck_flops = 2 * neck_macs
                        total_flops += neck_flops

                        # 打印颈部网络 FLOPs
                        print(f"颈部网络 FLOPs: {format_number(neck_flops)} ({neck_flops:,})")
                    except Exception as e:
                        print(f"计算颈部网络 FLOPs 失败: {e}")

                # 打印总 FLOPs 和参数量
                print(f"估计总 FLOPs (不包括检测头): {format_number(total_flops)} ({total_flops:,})")
                print(f"Params: {format_number(total_params)} ({total_params:,})")
                print(f"{split_line}")

            except Exception as thop_error:
                print(f"使用 thop 计算 FLOPs 失败: {thop_error}")

                # 如果两种方法都失败，只打印参数量
                print(f"{split_line}")
                print(f"Input shape: {input_shape}")
                print(f"FLOPs: 无法计算")
                print(f"Params: {format_number(total_params)} ({total_params:,})")
                print(f"{split_line}")

        # 打印模型结构
        print("\n模型结构:")
        print(f"{split_line}")
        print(model)
        print(f"{split_line}")

        # 打印各组件的参数量和占比
        print("\n各组件参数量:")
        print(f"{'-' * 60}")
        print(f"{'组件名称':<30} {'参数量':<15} {'占比':<10}")
        print(f"{'-' * 60}")

        # 分析主要组件
        for name, module in model.named_children():
            module_params = sum(p.numel() for p in module.parameters())
            percentage = 100.0 * module_params / total_params
            print(f"{name:<30} {format_number(module_params):<15} {percentage:.2f}%")

            # 分析子组件
            if hasattr(module, 'named_children'):
                for subname, submodule in module.named_children():
                    subparams = sum(p.numel() for p in submodule.parameters())
                    subpercentage = 100.0 * subparams / total_params
                    if subparams > total_params * 0.01:  # 只显示占比超过1%的子组件
                        print(f"  - {subname:<28} {format_number(subparams):<15} {subpercentage:.2f}%")

        print(f"{'-' * 60}")

    except Exception as e:
        print(f"分析模型时出错: {e}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()
